{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "[View the runnable example on GitHub](https://github.com/intel-analytics/BigDL/tree/main/python/nano/tutorial/notebook/inference/pytorch/accelerate_pytorch_inference_jit_ipex.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Accelerate PyTorch Inference using JIT/IPEX\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "* JIT: You can use `InferenceOptimizer.trace(..., accelerator=\"jit\")` API to enable the TorchScript acceleration for PyTorch inference.\n",
        "* IPEX: You can use `InferenceOptimizer.trace(...,use_ipex=True)` API to enable the IPEX (Intel® Extension for PyTorch*) acceleration for PyTorch inference.\n",
        "* JIT + IPEX: It is recommended to use JIT and IPEX together. You can user `InferenceOptimizer.trace(..., acclerator=\"jit\", use_ipex=True`) to enable both for PyTorch inference.\n",
        "\n",
        "All of the above accelerations only take a few lines to apply."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "nbsphinx": "hidden"
      },
      "source": [
        "To apply JIT/IPEX acceleration, you should install BigDL-Nano for PyTorch first："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "nbsphinx": "hidden"
      },
      "outputs": [],
      "source": [
        "!pip install --pre --upgrade bigdl-nano[pytorch] # install the nightly-built version\n",
        "!source bigdl-nano-init"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nbsphinx": "hidden"
      },
      "source": [
        "> 📝 **Note**\n",
        ">\n",
        "> We recommend to run the commands above, especially `source bigdl-nano-init` before jupyter kernel is started, or some of the optimizations may not take effect."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's take an [ResNet-18 model](https://pytorch.org/vision/main/models/generated/torchvision.models.resnet18.html) pretrained on ImageNet dataset as an example. First, we load the model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "from torchvision.models import resnet18\n",
        "\n",
        "model_ft = resnet18(pretrained=True)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "To accelerate inference using JIT, IPEX, or JIT together with IPEX, we need to import `InferenceOptimizer` first:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from bigdl.nano.pytorch import InferenceOptimizer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Accelerate Inference Using JIT Optimizer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "jit_model = InferenceOptimizer.trace(model_ft,\n",
        "                                     accelerator=\"jit\",\n",
        "                                     input_sample=torch.rand(1, 3, 224, 224))\n",
        "with InferenceOptimizer.get_context(jit_model):\n",
        "    x = torch.rand(2, 3, 224, 224)\n",
        "    y_hat = jit_model(x)\n",
        "    predictions = y_hat.argmax(dim=1)\n",
        "    print(predictions)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Accelerate Inference Using IPEX Optimizer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "ipex_model = InferenceOptimizer.trace(model_ft,\n",
        "                                      use_ipex=True)\n",
        "with InferenceOptimizer.get_context(ipex_model):\n",
        "    x = torch.rand(2, 3, 224, 224)\n",
        "    y_hat = ipex_model(x)\n",
        "    predictions = y_hat.argmax(dim=1)\n",
        "    print(predictions)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Accelerate Inference Using JIT + IPEX"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "jit_ipex_model = InferenceOptimizer.trace(model_ft,\n",
        "                                          accelerator=\"jit\",\n",
        "                                          use_ipex=True,\n",
        "                                          input_sample=torch.rand(1, 3, 224, 224))\n",
        "with InferenceOptimizer.get_context(jit_ipex_model):\n",
        "    x = torch.rand(2, 3, 224, 224)\n",
        "    y_hat = jit_ipex_model(x)\n",
        "    predictions = y_hat.argmax(dim=1)\n",
        "    print(predictions)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> 📝 **Note**\n",
        ">\n",
        "> `input_sample` is the parameter for accelerators to know the **shape** of the model input. So both the batch size and the specific values are not important to `input_sample`. If we want our test dataset to consist of images with $224 \\times 224$ pixels, we could use `torch.rand(1, 3, 224, 224)` for `input_sample` here.\n",
        ">\n",
        "> Please refer to [API documentation](https://bigdl.readthedocs.io/en/latest/doc/PythonAPI/Nano/pytorch.html#bigdl.nano.pytorch.InferenceOptimizer.trace) for more information on `InferenceOptimizer.trace`.\n",
        ">\n",
        "> Also note that for all Nano optimized models by `InferenceOptimizer.trace`, you need to wrap the inference steps with an automatic context manager `InferenceOptimizer.get_context(model=...)` provided by Nano. You could refer to [here](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Inference/PyTorch/pytorch_context_manager.html) for more detailed usage of the context manager."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> 📚 **Related Readings**\n",
        ">\n",
        "> - [How to install BigDL-Nano](https://bigdl.readthedocs.io/en/latest/doc/Nano/Overview/install.html)\n",
        "> - [How to enable automatic context management for PyTorch inference on Nano optimized models](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Inference/PyTorch/pytorch_context_manager.html)\n",
        "> - [How to save and load optimized JIT model](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Inference/PyTorch/pytorch_save_and_load_jit.html)\n",
        "> - [How to save and load optimized IPEX model](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Inference/PyTorch/pytorch_save_and_load_ipex.html)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "nano-pytorch",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.15 (default, Nov 24 2022, 21:12:53) \n[GCC 11.2.0]"
    },
    "orig_nbformat": 4,
    "vscode": {
      "interpreter": {
        "hash": "8a5edab282632443219e051e4ade2d1d5bbc671c781051bf1437897cbdfea0f1"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}
